#!/usr/bin/env python
"""
Celery启动脚本
用于启动Celery worker和beat调度器
"""

import os
import sys
import subprocess
from pathlib import Path

# 添加项目根目录到Python路径
BASE_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(BASE_DIR))

# 设置Django设置模块
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'demo0811.settings')


def start_celery_worker():
    """启动Celery worker"""
    print("启动Celery worker...")
    cmd = [
        'celery', '-A', 'demo0811', 'worker',
        '--loglevel=info',
        '--concurrency=4'
    ]
    return subprocess.Popen(cmd)


def start_celery_beat():
    """启动Celery beat调度器"""
    print("启动Celery beat调度器...")
    cmd = [
        'celery', '-A', 'demo0811', 'beat',
        '--loglevel=info',
        '--scheduler=django_celery_beat.schedulers:DatabaseScheduler'
    ]
    return subprocess.Popen(cmd)


if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='启动Celery服务')
    parser.add_argument('--worker-only', action='store_true', help='只启动worker')
    parser.add_argument('--beat-only', action='store_true', help='只启动beat调度器')
    
    args = parser.parse_args()
    
    processes = []
    
    try:
        if args.worker_only:
            processes.append(start_celery_worker())
        elif args.beat_only:
            processes.append(start_celery_beat())
        else:
            # 同时启动worker和beat
            processes.append(start_celery_worker())
            processes.append(start_celery_beat())
        
        print("Celery服务已启动，按Ctrl+C停止...")
        
        # 等待所有进程
        for process in processes:
            process.wait()
            
    except KeyboardInterrupt:
        print("\n正在停止Celery服务...")
        for process in processes:
            process.terminate()
        print("已停止所有Celery服务")
